{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "BASE_DIR = os.path.dirname(os.path.abspath(__file__))\n",
    "import numpy as np\n",
    "import torch\n",
    "import random\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from tools.my_dataset import RMBDataset\n",
    "# from tools.common_tools import set_seed, transform_invert\n",
    "\n",
    "def set_seed(seed=1):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "set_seed(1)  # 设置随机种子\n",
    "\n",
    "# 参数设置\n",
    "MAX_EPOCH = 10\n",
    "BATCH_SIZE = 1\n",
    "LR = 0.01\n",
    "log_interval = 10\n",
    "val_interval = 1"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# 在训练部分被引用\n",
    "# 传入两个参数img_和train_transform，返回PIL image，也就是可以直接plot将其格式化\n",
    "def transform_invert(img_, transform_train):\n",
    "    \"\"\"\n",
    "    将data 进行反transfrom操作\n",
    "    :param img_: tensor\n",
    "    :param transform_train: torchvision.transforms\n",
    "    :return: PIL image\n",
    "    \"\"\"\n",
    "\n",
    "    # 对normalize进行反操作\n",
    "    if 'Normalize' in str(transform_train):\n",
    "        norm_transform = list(filter(lambda x: isinstance(x, transforms.Normalize),\n",
    "                                     transform_train.transforms))\n",
    "        mean = torch.tensor(norm_transform[0].mean,\n",
    "                            dtype=img_.dtype,\n",
    "                            device=img_.device)\n",
    "        std = torch.tensor(norm_transform[0].std,\n",
    "                           dtype=img_.dtype,\n",
    "                           device=img_.device)\n",
    "        # normalize是减去均值除于方差，因此反操作就是乘于方差再加上均值\n",
    "        img_.mul_(std[:, None, None]).add_(mean[:, None, None])\n",
    "\n",
    "    # 通道变换，C*H*W --> H*W*C，将channel放到最后面\n",
    "    img_ = img_.transpose(0, 2).transpose(0, 1)\n",
    "\n",
    "    # 将0-1尺度上的数据转换到0-255\n",
    "    if 'ToTensor' in str(transform_train) or img_.max() < 1:\n",
    "        img_ = img_.detach().numpy() * 255\n",
    "\n",
    " # 将np_array的形式转换成PIL image\n",
    "    # 判断channel是3通道还是1通道，分别转换成RGB彩色图像和灰度图像\n",
    "    if img_.shape[2] == 3:\n",
    "        img_ = Image.fromarray(img_.astype('uint8')).convert('RGB')\n",
    "    elif img_.shape[2] == 1:\n",
    "        img_ = Image.fromarray(img_.astype('uint8').squeeze())\n",
    "    else:\n",
    "        raise Exception(\"Invalid img shape, expected 1 or 3 in axis 2, but got {}!\"\n",
    "                        .format(img_.shape[2]) )\n",
    " # 返回图像就可以对图像进行plot，对图像进行可视化\n",
    "    return img_"
   ],
   "id": "cfa9609a5dade39f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "895f20f018e50a7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
